import warnings
warnings.filterwarnings('ignore')
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
tf.compat.v1.logging.set_verbosity(40)
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense,Input
from tensorflow.keras.applications import MobileNet
from tensorflow.keras.losses import SparseCategoricalCrossentropy

from sklearn.model_selection import train_test_split
import numpy as np
import matplotlib.pyplot as plt
import random
import cv2


def readData(path):
    x_data = []
    y_data = []
    for i,label_name in enumerate(os.listdir(path)):
        label_path = os.path.join(path,label_name)
        for img_name in os.listdir(label_path):
            img_path = os.path.join(label_path,img_name)
            img = cv2.cvtColor(cv2.imread(img_path),cv2.COLOR_BGR2RGB)
            img = cv2.resize(img,(128,128))/255
            x_data.append(img)
            y_data.append(i)
    return np.array(x_data),np.array(y_data)

def MyCNN(input_shape,num_classes):
    inputs = Input(input_shape)
    pre_train_model = MobileNet(input_shape=input_shape,include_top=False,pooling='avg')
    pre_train_model.trainable = False
    x = pre_train_model(inputs)
    outputs = Dense(num_classes,activation='softmax')(x)
    model = Model(inputs,outputs)
    return model

if __name__ == '__main__':
    data_path = r'E:\DATA\Gesture_Recognition'
    classes_name = os.listdir(data_path)
    x_data,y_data = readData(data_path)

    x_train,x_test,y_train,y_test = train_test_split(x_data,y_data,train_size=0.8)
    x_val,x_test,y_val,y_test = train_test_split(x_test,y_test,train_size=0.5)

    model = MyCNN(x_data.shape[1:],len(classes_name))
    model.summary()

    model.compile(loss=SparseCategoricalCrossentropy(),optimizer='adam',metrics=['accuracy'])
    model.fit(x_train,y_train,epochs=20,batch_size=8,validation_data=(x_val,y_val))

    score = model.evaluate(x_test,y_test,verbose=0)
    print(f'测试集损失值:{score[0]:.3f}')
    print(f'测试集准确率:{score[1]:.3f}')

    plt.figure(figsize=(9,9))
    for i in range(9):
        plt.subplot(3,3,i+1)
        r = random.randint(0,len(x_test)-1)
        plt.imshow(x_test[r])
        pred = np.argmax(model.predict(x_test[r:r+1]),1)
        pred_label = classes_name[pred[0]]
        true_label = classes_name[y_test[r]]
        plt.title(pred_label,c='k' if pred_label==true_label else 'r')
        plt.axis('off')
    plt.show()